package cn.baiyang.rocketmq.spring.core;

import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.CHANGED_HOLDER;
import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.DEFAULT_QPS;
import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.INITIAL_TOKENS;
import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.PARAM_HOLDER;
import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.TOPICTAG_LIMIT;
import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.TOPICTAG_LIMIT_ENABLED;
import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.TOPIC_LIMIT;
import static cn.baiyang.rocketmq.spring.support.ratelimit.RateLimiterParamsHolder.TOPIC_LIMIT_ENABLED;

import java.time.Duration;
import java.util.Arrays;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.apache.rocketmq.client.apis.ClientException;
import org.apache.rocketmq.client.apis.message.Message;
import org.apache.rocketmq.client.apis.message.MessageBuilder;
import org.apache.rocketmq.client.apis.message.MessageId;
import org.apache.rocketmq.client.apis.producer.Producer;
import org.apache.rocketmq.client.apis.producer.SendReceipt;
import org.apache.rocketmq.client.apis.producer.Transaction;
import org.apache.rocketmq.client.java.message.MessageBuilderImpl;
import org.apache.rocketmq.shaded.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.env.Environment;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.BucketConfiguration;
import io.github.bucket4j.ConfigurationBuilder;
import io.github.bucket4j.TokensInheritanceStrategy;
import io.github.bucket4j.distributed.proxy.ProxyManager;

/**
 * RocketMQ Message Sender
 *
 * @author hongzhu
 * @version V1.0
 * @since 2023-03-13 10:00
 */
public abstract class RocketMQSender {
    private static final Logger LOGGER = LoggerFactory.getLogger(RocketMQSender.class);

    private static final ExecutorService CALLBACK_EXECUTOR_SERVICE = Executors.newCachedThreadPool();

    protected Environment environment;

    protected ProxyManager<String> proxyManager;

    protected Producer producer;

    protected abstract void setEnvironment(Environment environment);

    protected abstract void setProxyManager(ProxyManager<String> proxyManager);

    protected abstract void setProducer(Producer producer);

    public Producer getProducer() {
        return producer;
    }

    /**
     * 发送成功后返回message id
     * @param message {@link org.apache.rocketmq.client.apis.message.Message}
     * @param transaction {@link org.apache.rocketmq.client.apis.producer.Transaction}
     * @throws ClientException sending exception
     * @return Message Id String
     */
    String send(Message message, Transaction transaction) throws ClientException {
        SendReceipt sendReceipt = producer.send(message, transaction);
        MessageId messageId = null == sendReceipt ? null : sendReceipt.getMessageId();
        LOGGER.debug("Send transaction message successfully, messageId={}", null == messageId ? "" : messageId);
        return null == messageId ? "" : messageId.toString();
    }

    Message buildMessage(String topic, String tag, String messageGroup,
        byte[] body, Long deliveryTimestamp, Map<String, String> properties, String[] keys) {
        MessageBuilder messageBuilder = new MessageBuilderImpl();
        messageBuilder = messageBuilder.setTopic(topic);
        if (StringUtils.isNotBlank(tag)) { // 设置消息Tag，用于消费端根据指定Tag过滤消息
            messageBuilder = messageBuilder.setTag(tag);
        }
        if (null != deliveryTimestamp && StringUtils.isNotBlank(messageGroup)) {
            throw new RuntimeException("请检查消息类型，不能同时发送顺序消息和延迟消息");
        }
        if (StringUtils.isNotBlank(messageGroup)) { // 设置message group来保证顺序消息
            messageBuilder = messageBuilder.setMessageGroup(messageGroup);
        }
        if (null != deliveryTimestamp) { // 设置延迟投递消息的时间，用于延时/定时消息
            messageBuilder = messageBuilder.setDeliveryTimestamp(deliveryTimestamp);
        }
        if (null != keys && keys.length > 0) { // 设置消息索引键，可根据关键字精确查找某条消息
            if (!Arrays.asList(keys).contains(null)){
                messageBuilder = messageBuilder.setKeys(keys);
            }
        }
        if (null != properties && !properties.isEmpty()) { // 设置附加属性，可用于SQL92类型的消息过滤
            for (Map.Entry<String, String> e : properties.entrySet()) {
                messageBuilder = messageBuilder.addProperty(e.getKey(), e.getValue());
            }
        }
        return messageBuilder.setBody(body).build();
    }

    /**
     * 发送成功后返回message id, 或者失败返回空字符
     * @param message {@link org.apache.rocketmq.client.apis.message.Message}
     * @throws ClientException sending exception
     * @return Message Id String
     */
    String send(Message message) throws ClientException {
        if (null != proxyManager) {
            ratelimit(message.getTopic(), message.getTag());
        }

        // 发送消息，需要关注发送结果，并捕获失败等异常。
        SendReceipt sendReceipt = producer.send(message);
        MessageId messageId = null == sendReceipt ? null : sendReceipt.getMessageId();
        LOGGER.debug("Send message successfully, messageId={}", null == messageId ? "" : messageId);
        return null == messageId ? "" : messageId.toString();
    }

    /**
     * 发送成功后返回message id, 或者失败返回空字符<br/>
     * 异步发送时立即返回为空字符
     * @param message {@link org.apache.rocketmq.client.apis.message.Message}
     * @return Message Id String
     */
    String sendAsync(Message message) {
        if (null != proxyManager) {
            ratelimit(message.getTopic(), message.getTag());
        }

        final String[] messageId = new String[]{""};
        CompletableFuture<SendReceipt> completableFuture = producer.sendAsync(message);
        completableFuture.whenCompleteAsync((sendReceipt, throwable) -> {
            if (null != throwable) {
                LOGGER.error("Failed to send async message on complete", throwable);
                return;
            }
            messageId[0] = sendReceipt.getMessageId().toString();
            LOGGER.debug("Send async message successfully, messageId={}", messageId[0]);
        }, CALLBACK_EXECUTOR_SERVICE);
        return messageId[0];
    }

    private void ratelimit(String topic, Optional<String> tag) {
        String qpsName = tag.isPresent()
            ? String.format(TOPICTAG_LIMIT, topic, tag.get()) : String.format(TOPIC_LIMIT, topic);
        String enabledName = tag.isPresent()
            ? String.format(TOPICTAG_LIMIT_ENABLED, topic, tag.get()) : String.format(TOPIC_LIMIT_ENABLED, topic);
        boolean limitEnabled = false;
        Bandwidth bandwidth = null;
        if (tag.isPresent()) { // (Topic + Tag)维度优先级高于Topic维度
            limitEnabled = getLimitEnabled(enabledName);
            if (limitEnabled) {
                try { // 如果有(Topic + Tag)维度自定义配置
                    environment.resolveRequiredPlaceholders(qpsName);
                    bandwidth = buildBandwidth(qpsName);
                } catch (Exception e) { // 否则检查Topic维度自定义配置，默认行为
                    LOGGER.debug("No (Topic + Tag)={} Dimension RateLimit Configured, Use Default Topic Dimension RateLimit.", qpsName, e);
                    qpsName = String.format(TOPIC_LIMIT, topic);
                    bandwidth = buildBandwidth(qpsName);
                }
            }
        } else { // 检查Topic维度自定义配置
            limitEnabled = getLimitEnabled(enabledName);
            if (limitEnabled) { // 若是没有自定义配置，则填充默认的50qps限流
                bandwidth = buildBandwidth(qpsName);
            }
        }

        if (limitEnabled) {// 只要自定义限流配置enabled不为false，都要做限流；包括默认没有配置enabled项
            ConfigurationBuilder configBuilder = BucketConfiguration.builder();
            BucketConfiguration bucketConfiguration = configBuilder.addLimit(bandwidth).build();
            Bucket bucket;
            try {
                if (CHANGED_HOLDER.getOrDefault(qpsName, false)) {
                    /*
                     * only the rate limiter parameters is configured in config center and changed, then we go here.
                     */
                    bucket = proxyManager.builder()
                        .withImplicitConfigurationReplacement(System.currentTimeMillis(), TokensInheritanceStrategy.PROPORTIONALLY)
                        .build(qpsName, () -> {
                            return bucketConfiguration;
                        });
                    CHANGED_HOLDER.put(qpsName, false);
                } else {
                    bucket = proxyManager.builder().build(qpsName, () -> {
                        return bucketConfiguration;
                    });
                }
                // 超过限流速度进行等待
                bucket.asBlocking().consume(1);
            } catch (Exception e) {
                LOGGER.error("RateLimit Failed On {}, Interrupted Exception Occurred On Bucket4j.", qpsName, e);
            }
        }
    }

    private Bandwidth buildBandwidth(String qpsName) {
        String qpsStr = environment.resolvePlaceholders(qpsName);
        long qps = getQps(qpsStr);
        Long oldQps = PARAM_HOLDER.get(qpsName);
        if (null == oldQps) {
            PARAM_HOLDER.put(qpsName, qps);
        } else {
            if (oldQps != qps) {
                LOGGER.info("RateLimiter value of {} is changed from {} to {}", qpsName, oldQps, qps);
                CHANGED_HOLDER.put(qpsName, true);
                PARAM_HOLDER.put(qpsName, qps);
            }
        }
        return Bandwidth.builder()
            .capacity(qps).refillGreedy(qps, Duration.ofSeconds(1)).initialTokens(INITIAL_TOKENS)
            .id(qpsName).build();
    }

    private long getQps(String qpsStr) {
        long qps;
        try {
            qps = Long.parseLong(qpsStr);
        } catch (Exception e) {
            LOGGER.debug("qps is {}, not configured or not correct. Fallback to default qps", qpsStr, e);
            qps = DEFAULT_QPS;
        }
        qps = qps <= 0 ? DEFAULT_QPS : qps;
        return qps;
    }

    private Boolean getLimitEnabled(String enabledName) {
        try {
            String enabled = environment.resolveRequiredPlaceholders(enabledName);
            return Boolean.parseBoolean(enabled);
        } catch (Exception e) {
            LOGGER.debug("{} Not Configured, Use Default: true", enabledName, e);
            return true;
        }
    }

}
